import numpy as np
import matplotlib.pyplot as plt

# -----------------------------------------------------------------------------
# ECUACION CONVECCION 1D ------------------------------------------------------
# -----------------------------------------------------------------------------
def FTCS_conveccion(N, tiempo, delta_t, delta_x, u, init_func, cf_func):
    
    C = u * delta_t / delta_x
    print(f"C ={C:.2f}")
    
    # Calculamos iteraciones
    iteraciones = int(tiempo / delta_t)
    
    # Condiciones iniciales (dadas por nuestra funcion)
    T = init_func(N)
    Tnew = np.zeros(N+1)    # Y formamos un array identico a T pero con todo ceros
    
    # Ploteamos los valores del primer tiempo inicial
    plt.figure(1)
    plt.plot(T)
    
    for n in range(iteraciones):    # Opera tiempo a tiempo (iteracion a iteracion)
        
        for i in range(1, N):   # Calculamos Tnew de cada particula menos las dos de los lados
            Tnew[i] = T[i] - (1/2)*C*(T[i+1] - T[i-1])
        
        # Condiciones de frontera externas
        Tnew = cf_func(Tnew, n, delta_t)    # Usamos la funcion, que puede necesitar n o dt (como sen())
        
        T[:] = Tnew[:]  # Actualiza datos antes de pasar a siguiente interaccion
        
        if n % 10 == 0:    # Mostrar como fotogramas de momentos, cada 10 iteraciones una foto
            #plt.figure(1)
            plt.plot(T)
            plt.pause(0.01)
            plt.xlabel("i-indice")
            plt.ylabel("Temp")
    plt.show()
    
    return T

# Funcion donde todos los valores iniciales son 0
def init_None(N):
    T = np.zeros(N+1)
    return T

# Funcion donde todos los valores iniciales son random entre 0 y 10
def init_random(N):
    T = np.random.rand(N+1)*10  # Valores entre 0 y 10
    return T

# Las fronteras son fijas, focos fijos en los lados (0 en la izquierda, 10 en la derecha)
def dirichlet(T, n, dt):
    T[0] = 0
    T[-1] = 10
    return T

# Funcion rara donde la frontera izquierda es siempre 0 y la derecha va oscilando
def seno(T, n, dt):
    t = n * dt
    T[0] = 0
    T[-1] = np.sin(0.5 * t)
    return T

# Condiciones de frontera de flujo nulo, como si fuera una piscina, los bordes van moviendose
def flujo_nulo(T, n, dt):
    T[0] = T[1]
    T[-1] = T[-2]
    return T

# EJERCICIO 1 -----------------------------------------------------------------

FTCS_conveccion(N=20, tiempo=10, delta_t=0.01, delta_x=0.1, u=1, init_func=init_None, cf_func=dirichlet)


# EJERCICIO 2 -----------------------------------------------------------------

def upwind(N, tiempo, delta_t, delta_x, u, init_func, cf_func):
    
    C = u * delta_t / delta_x
    print(f"C ={C:.2f}")
    if C>1:
        print("r inestable. Aumenta delta_x o reduce delta_t")
    elif C ==1:
        print("Solución exacta.")
    else:
        pass
    
    # Calculamos iteraciones
    iteraciones = int(tiempo / delta_t)
    
    # Condiciones iniciales (dadas por nuestra funcion)
    T = init_func(N)
    Tnew = np.zeros(N+1)    # Y formamos un array identico a T pero con todo ceros
    
    # Ploteamos los valores del primer tiempo inicial
    plt.figure(1)
    plt.plot(T)
    
    for n in range(iteraciones):    # Opera tiempo a tiempo (iteracion a iteracion)
        
        for i in range(1, N):   # Calculamos Tnew de cada particula menos las dos de los lados
            Tnew[i] = T[i]*(1-C) + C*T[i-1]
            
        
        # Condiciones de frontera externas
        Tnew = cf_func(Tnew, n, delta_t)    # Usamos la funcion, que puede necesitar n o dt (como sen())
        
        T[:] = Tnew[:]  # Actualiza datos antes de pasar a siguiente interaccion
        
        if n % 1 == 0:    # Mostrar como fotogramas de momentos, cada 10 iteraciones una foto
            #plt.figure(1)
            plt.plot(T)
            plt.pause(0.01)
            plt.xlabel("i-indice")
            plt.ylabel("Temp")
    plt.show()
    
    return T

def pulso_rio(N):
    T = np.zeros(N+1)
    T[0]=1
    return T

def rio_None(T, n, dt):
    return T

upwind(N=20, tiempo=5, delta_t=0.1, delta_x=0.1, u=1, init_func=pulso_rio, cf_func=rio_None)


# EJERCICIO 3 -----------------------------------------------------------------

def DF_conveccion(N, tiempo, delta_t, delta_x, u, init_func, cf_func):
    
    C = u * delta_t / delta_x
    print(f"C ={C:.2f}")
    if C>=1:
        print("r inestable. Aumenta delta_x o reduce delta_t")
    else:
        pass
    
    # Calculamos iteraciones
    iteraciones = int(tiempo / delta_t)
    
    # Condiciones iniciales (dadas por nuestra funcion)
    T = init_func(N)
    Told = np.copy(T)
    Tnew = np.zeros(N+1)
    # Para salvar situacion, avanzar un tiempo aplicando FTCS que se aproxima a nuestro modelo como arranque
    T_primer_paso = np.copy(T)  # Primer paso de FTCS aproximado y luego ya lo normal y Tn dirige
    for i in range(1, N):
        T_primer_paso[i] = T[i] - (1/2)*C*(T[i+1] - T[i-1])
        # Tnew[i] = T[i]*(1-C) + C*T[i-1]
    
    T_primer_paso = cf_func(T_primer_paso, 0, delta_t)   # COndiciones de frontera para el momento init, n=0
    T = np.copy(T_primer_paso)
    
    # Ploteamos los valores del primer tiempo inicial
    plt.figure(1)
    plt.plot(T)
    
    for n in range(iteraciones):    # Opera tiempo a tiempo (iteracion a iteracion)
        
        for i in range(1, N):   # Calculamos Tnew de cada particula menos las dos de los lados
            Tnew[i] = Told[i] - C*(T[i+1] - T[i-1])
            
        
        # Condiciones de frontera externas
        Tnew = cf_func(Tnew, n, delta_t)    # Usamos la funcion, que puede necesitar n o dt (como sen())
        
        Told[:] = T[:]
        T[:] = Tnew[:]  # Actualiza datos antes de pasar a siguiente interaccion
        
        if n % 1 == 0:    # Mostrar como fotogramas de momentos, cada 10 iteraciones una foto
            #plt.figure(1)
            plt.plot(T)
            plt.pause(0.01)
            plt.xlabel("i-indice")
            plt.ylabel("Temp")
    plt.show()
    
    return T

# Funcion de un "pulso" en el pirmer punto
def pulso_rio(N):
    T = np.zeros(N+1)
    T[0]=1
    return T
# Sin condiciones de frontera
def rio_None(T, n, dt):
    return T

DF_conveccion(N=20, tiempo=3, delta_t=0.1, delta_x=0.1, u=1/2, init_func=pulso_rio, cf_func=rio_None)


# -----------------------------------------------------------------------------
# ECUACION TRANSPORTE COMPLETA ------------------------------------------------
# -----------------------------------------------------------------------------

# EJERCICIO 6 -----------------------------------------------------------------

def transporte_FTCS(N, tiempo, delta_t, delta_x, u, alpha, init_func, cf_func):
    
    C = u * delta_t / delta_x
    print(f"C = {C:.2f}")
    s = alpha * delta_t / delta_x**2
    print(f"s = {s:.2f}")
    print("\nEstable si 0<=C^2<=2s<=1")
    if C**2 > 1 or 2*s > 1 or C**2 > 2*s:
        print("INESTABLE")
    else:
        print("ESTABLE")
    
    # Calculamos iteraciones
    iteraciones = int(tiempo / delta_t)
    
    # Condiciones iniciales (dadas por nuestra funcion)
    T = init_func(N)
    Tnew = np.zeros(N+1)    # Y formamos un array identico a T pero con todo ceros
    
    # Ploteamos los valores del primer tiempo inicial
    plt.figure(1)
    plt.plot(T)
    
    for n in range(iteraciones):    # Opera tiempo a tiempo (iteracion a iteracion)
        
        for i in range(1, N):   # Calculamos Tnew de cada particula menos las dos de los lados
            Tnew[i] = s * (T[i-1] + T[i+1] - 2*T[i]) + (1/2) * C * (T[i-1] - T[i+1]) + T[i]
        
        # Condiciones de frontera externas
        Tnew = cf_func(Tnew, n, delta_t)    # Usamos la funcion, que puede necesitar n o dt (como sen())
        
        T[:] = Tnew[:]  # Actualiza datos antes de pasar a siguiente interaccion
        
        if n % 10 == 0:    # Mostrar como fotogramas de momentos, cada 10 iteraciones una foto
            #plt.figure(1)
            plt.plot(T)
            plt.pause(0.01)
            plt.xlabel("i-indice")
            plt.ylabel("Temp")
    plt.show()
    
    return T

transporte_FTCS(N=20, tiempo=3, delta_t=0.01, delta_x=0.1, u=1, alpha = 1/3, init_func=pulso_rio, cf_func=rio_None)

# EJERCICIO 7 -----------------------------------------------------------------

def transporte_upwind(N, tiempo, delta_t, delta_x, u, alpha, init_func, cf_func):
    
    C = u * delta_t / delta_x
    print(f"C = {C:.2f}")
    s = alpha * delta_t / delta_x**2
    print(f"s = {s:.2f}")
    print("\nEstable si 0<=C^2<=2s<=1")
    if C**2 > 1 or 2*s > 1 or C**2 > 2*s:
        print("INESTABLE")
    else:
        print("ESTABLE")
    
    # Calculamos iteraciones
    iteraciones = int(tiempo / delta_t)
    
    # Condiciones iniciales (dadas por nuestra funcion)
    T = init_func(N)
    Tnew = np.zeros(N+1)    # Y formamos un array identico a T pero con todo ceros
    
    # Ploteamos los valores del primer tiempo inicial
    plt.figure(1)
    plt.plot(T)
    
    for n in range(iteraciones):    # Opera tiempo a tiempo (iteracion a iteracion)
        
        for i in range(1, N):   # Calculamos Tnew de cada particula menos las dos de los lados
            Tnew[i] = T[i] + C * (T[i-1] - T[i]) + s * (T[i-1] + T[i+1] - 2* T[i])
        
        # Condiciones de frontera externas
        Tnew = cf_func(Tnew, n, delta_t)    # Usamos la funcion, que puede necesitar n o dt (como sen())
        
        T[:] = Tnew[:]  # Actualiza datos antes de pasar a siguiente interaccion
        
        if n % 10 == 0:    # Mostrar como fotogramas de momentos, cada 10 iteraciones una foto
            #plt.figure(1)
            plt.plot(T)
            plt.pause(0.01)
            plt.xlabel("i-indice")
            plt.ylabel("Temp")
    
    plt.show()
    
    return T

transporte_upwind(N=20, tiempo=3, delta_t=0.01, delta_x=0.1, u=1, alpha = 1/3, init_func=pulso_rio, cf_func=rio_None)

# EJERCICIO 8 -----------------------------------------------------------------

def transporte_DF(N, tiempo, delta_t, delta_x, u, alpha, init_func, cf_func):
    
    C = u * delta_t / delta_x
    print(f"C = {C:.2f}")
    s = alpha * delta_t / delta_x**2
    print(f"s = {s:.2f}")
    print("\nEstable si 0<=C^2<=2s<=1")
    if C**2 > 1 or 2*s > 1 or C**2 > 2*s:
        print("INESTABLE")
    else:
        print("ESTABLE")
    
    # Calculamos iteraciones
    iteraciones = int(tiempo / delta_t)
    
    # Condiciones iniciales (dadas por nuestra funcion)
    T = init_func(N)
    Told = np.copy(T)
    Tnew = np.zeros(N+1)
    T_primer_paso = np.copy(T)
    for i in range(1, N):   # FTCS
        T_primer_paso[i] = s * (T[i-1] + T[i+1] - 2*T[i]) + (1/2) * C * (T[i-1] - T[i+1]) + T[i]
    
    T_primer_paso = cf_func(T_primer_paso, 0, delta_t)   # COndiciones de frontera para el momento init, n=0
    T = np.copy(T_primer_paso)
    
    # Ploteamos los valores del primer tiempo inicial
    plt.figure(1)
    plt.plot(T)
    
    for n in range(iteraciones):    # Opera tiempo a tiempo (iteracion a iteracion)
        
        for i in range(1, N):   # Calculamos Tnew de cada particula menos las dos de los lados
            Tnew[i] = (1/(1 + 2*s))*(Told[i] + C*(T[i-1] - T[i+1]) + 2*s*(T[i-1] + T[i+1] - Told[i]))
        
        # Condiciones de frontera externas
        Tnew = cf_func(Tnew, n, delta_t)    # Usamos la funcion, que puede necesitar n o dt (como sen())
        
        T[:] = Tnew[:]  # Actualiza datos antes de pasar a siguiente interaccion
        
        if n % 10 == 0:    # Mostrar como fotogramas de momentos, cada 10 iteraciones una foto
            #plt.figure(1)
            plt.plot(T)
            plt.pause(0.01)
            plt.xlabel("i-indice")
            plt.ylabel("Temp")
    
    plt.show()
    
    return T

transporte_upwind(N=20, tiempo=3, delta_t=0.01, delta_x=0.1, u=1, alpha = 1/3, init_func=pulso_rio, cf_func=rio_None)

# -----------------------------------------------------------------------------
# ECUACIONES MULTIDIMENSIONALES -----------------------------------------------
# -----------------------------------------------------------------------------

# EJERCICIO 12 ----------------------------------------------------------------